import sys
import time

import cv2
import numpy as np
from PyQt5.QtCore import Qt, QRect, QSettings, QThread, pyqtSignal
from PyQt5.QtGui import QPainter, QColor, QPen, QFontDatabase, QFont
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QLabel,
                             QVBoxLayout, QPushButton, QHBoxLayout, QTextEdit)
from openai import OpenAI
from paddleocr import PaddleOCR

client = OpenAI(api_key="这里写你的API秘钥", base_url="https://api.deepseek.com")


class ResizableFrame(QWidget):
    MARGIN = 8  # 边缘检测阈值

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.settings = QSettings("MyApp", "ScreenTranslator")
        self.load_position()
        self.mouse_press_pos = None
        self.resize_direction = None
        self.setMouseTracking(True)

    def load_position(self):
        rect = self.settings.value("frame_geometry", QRect(100, 100, 400, 300))
        self.setGeometry(rect)

    def save_position(self):
        self.settings.setValue("frame_geometry", self.geometry())

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setPen(QPen(QColor(0, 255, 0), 2, Qt.DashLine))
        painter.setBrush(QColor(255, 255, 255, 30))
        painter.drawRect(self.rect())

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.mouse_press_pos = event.pos()
            self.resize_direction = self.get_resize_direction(event.pos())

    def mouseMoveEvent(self, event):
        pos = event.pos()
        direction = self.get_resize_direction(pos)
        cursor_map = {
            "top": Qt.SizeVerCursor,
            "bottom": Qt.SizeVerCursor,
            "left": Qt.SizeHorCursor,
            "right": Qt.SizeHorCursor,
            "top-left": Qt.SizeFDiagCursor,
            "top-right": Qt.SizeBDiagCursor,
            "bottom-left": Qt.SizeBDiagCursor,
            "bottom-right": Qt.SizeFDiagCursor
        }
        self.setCursor(cursor_map.get(direction, Qt.ArrowCursor))
        if self.mouse_press_pos and self.resize_direction:
            rect = self.geometry()
            delta = event.pos() - self.mouse_press_pos

            if 'left' in self.resize_direction:
                rect.setLeft(rect.left() + delta.x())
            if 'right' in self.resize_direction:
                rect.setRight(rect.right() + delta.x())
            if 'top' in self.resize_direction:
                rect.setTop(rect.top() + delta.y())
            if 'bottom' in self.resize_direction:
                rect.setBottom(rect.bottom() + delta.y())
            self.setGeometry(rect.normalized())
            self.mouse_press_pos = event.pos()  # 更新鼠标按下位置
        elif event.buttons() == Qt.LeftButton:
            self.move(self.pos() + event.pos() - self.mouse_press_pos)

    def mouseReleaseEvent(self, event):
        self.mouse_press_pos = None
        self.resize_direction = None
        self.save_position()

    def get_resize_direction(self, pos):
        rect = self.rect()
        x, y = pos.x(), pos.y()
        directions = []
        if x <= self.MARGIN:
            directions.append("left")
        elif x >= rect.width() - self.MARGIN:
            directions.append("right")
        if y <= self.MARGIN:
            directions.append("top")
        elif y >= rect.height() - self.MARGIN:
            directions.append("bottom")
        return "-".join(directions) if directions else None


class TranslatorApp(QMainWindow):
    class TranslationThread(QThread):
        update_result_signal = pyqtSignal(str)

        def __init__(self, user_input, parent=None):
            super().__init__(parent)
            self.user_input = user_input

        def run(self):
            try:
                response = client.chat.completions.create(
                    model="deepseek-chat",
                    messages=[{
                        "role": "system",
                        "content": "你是一个翻译专家，你能将日语或英语翻译成中文，下文是我使用本地OCR识别出的文本，请将这些文本翻译成中文，由于是OCR识别出的，所以有可能会出现错误，你在翻译的时候请尽量考虑到这点，另外，这些文本大多是游戏内的对话，请按照原本的风格翻译出来，最后，你只需要翻译，不需要对翻译内容或翻译思路做任何解释，不需要说任何多余的话"
                    },
                        {"role": "user", "content": self.user_input}],
                    stream=False,
                    temperature=1.3
                )
                translated_text = response.choices[0].message.content
                self.update_result_signal.emit(f"翻译结果：\n{translated_text}")
            except Exception as e:
                self.update_result_signal.emit(f"翻译错误：{str(e)}")

    def __init__(self):
        super().__init__()
        self.ocr = PaddleOCR(use_angle_cls=True, lang='en')  # 使用 PaddleOCR
        self.init_ui()

    def init_ui(self):
        # 加载“萝莉体第二版”字体
        font_id = QFontDatabase.addApplicationFont(
            "C:\\Users\\Lenovo\\AppData\\Local\\Microsoft\\Windows\\Fonts\\萝莉体.ttc")
        font_family = QFontDatabase.applicationFontFamilies(font_id)[0]
        app_font = QFont(font_family)
        self.setFont(app_font)  # 设置全局字体
        self.setWindowTitle("实时翻译工具")
        self.setGeometry(300, 300, 400, 300)
        main_widget = QWidget()
        layout = QVBoxLayout()
        control_layout = QHBoxLayout()
        self.btn_show = QPushButton("显示选区框")
        self.btn_show.clicked.connect(self.toggle_frame)
        # 添加 OCR 按钮
        self.btn_ocr = QPushButton("OCR识别")
        self.btn_ocr.clicked.connect(self.perform_ocr)
        # 添加 翻译 按钮
        self.btn_translate = QPushButton("翻译")
        self.btn_translate.clicked.connect(self.start_translation)
        self.btn_translate.setEnabled(False)
        control_layout.addWidget(self.btn_show)
        control_layout.addStretch()
        control_layout.addWidget(self.btn_ocr)
        control_layout.addWidget(self.btn_translate)
        # 添加 QTextEdit 组件用于显示和编辑 OCR 结果
        self.text_edit = QTextEdit()
        self.text_edit.setPlaceholderText("OCR识别出的文本将在此显示，您可以修改它。")
        self.result_label = QLabel("等待首次翻译...")
        self.result_label.setWordWrap(True)
        self.result_label.setStyleSheet("""
            QLabel {
                background: #f0f0f0;
                padding: 10px;
                border-radius: 5px;
                min-height: 120px;
                font-size: 14px;
                color: #333;
            }
        """)

        # 使用“萝莉体第二版”字体
        self.result_label.setFont(app_font)
        layout.addLayout(control_layout)
        layout.addWidget(self.text_edit)
        layout.addWidget(self.result_label)
        main_widget.setLayout(layout)
        self.setCentralWidget(main_widget)
        self.frame = ResizableFrame()
        self.frame.hide()
        # 设置主窗口的样式
        self.setStyleSheet("""
            QMainWindow {
                background: linear-gradient(to bottom, #4CAF50, #81C784);
                color: #fff;
                font-family: Arial, sans-serif;
            }
            QPushButton {
                background: #388E3C;
                color: white;
                padding: 10px;
                border-radius: 5px;
                font-size: 14px;
                border: none;
                box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1);
            }
            QPushButton:hover {
                background: #66BB6A;
            }
            QTextEdit {
                background: #fff;
                border-radius: 5px;
                padding: 10px;
                font-size: 14px;
                color: #333;
            }
        """)

    def toggle_frame(self):
        if self.frame.isVisible():
            self.frame.hide()
            self.btn_show.setText("显示选区框")
            self.btn_translate.setEnabled(False)
        else:
            self.frame.show()
            self.btn_show.setText("隐藏选区框")
            self.btn_translate.setEnabled(True)

    def get_screenshot(self):
        geo = self.frame.geometry()
        screen = QApplication.primaryScreen()
        qimg = screen.grabWindow(0, geo.x(), geo.y(), geo.width(), geo.height()).toImage()
        ptr = qimg.bits()
        ptr.setsize(qimg.byteCount())
        arr = np.array(ptr).reshape(qimg.height(), qimg.width(), 4)
        return cv2.cvtColor(arr, cv2.COLOR_RGBA2RGB)

    def perform_ocr(self):
        try:
            if not self.frame.isVisible():
                self.result_label.setText("请先显示选区框并选择区域")
                return
            img = self.get_screenshot()
            ocr_time = time.time()  # 记录OCR识别开始时间
            results = self.ocr.ocr(img, cls=True)
            ocr_duration = time.time() - ocr_time  # OCR识别耗时
            ocr_text = "\n".join([res[1][0] for res in results[0] if res[1]]) if results else "未识别到文本"
            # 将 OCR 结果显示到文本框
            self.text_edit.setPlainText(ocr_text)
            self.result_label.setText(f"OCR识别完成，您可以修改文本。")
        except Exception as e:
            self.result_label.setText(f"OCR识别错误：{str(e)}")

    def start_translation(self):
        user_input = self.text_edit.toPlainText()  # 获取用户修改后的文本
        # 确保文本不为空
        if not user_input.strip():
            self.result_label.setText("请输入需要翻译的文本。")
            return
        # 显示等待信息
        self.result_label.setText("正在翻译，请稍候...")
        # 启动翻译线程
        self.translation_thread = self.TranslationThread(user_input)
        self.translation_thread.update_result_signal.connect(self.update_translation_result)
        self.translation_thread.start()

    def update_translation_result(self, result_text):
        self.result_label.setText(result_text)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = TranslatorApp()
    window.show()
    sys.exit(app.exec_())
